package com.sunzm.spark

import java.util.Date

import com.alibaba.fastjson.JSON
import com.sunzm.spark.SparkSortDemo.SortOrder
import org.apache.commons.lang3.StringUtils
import org.apache.commons.lang3.time.DateUtils
import org.apache.log4j.{Level, Logger}
import org.apache.spark.Partitioner
import org.apache.spark.rdd.RDD
import org.apache.spark.sql.{DataFrame, SparkSession}
import org.apache.spark.util.Utils

/**
 *
 * 使用spark对文件进行排序
 *
 * @author Administrator
 * @version 1.0
 * @date 2021-06-08 21:02
 */
object SparkSortDemo {
  def main(args: Array[String]): Unit = {
    Logger.getLogger("org").setLevel(Level.WARN)

    val spark: SparkSession = SparkSession
      .builder()
      .appName(s"${this.getClass.getSimpleName.stripSuffix("$")}")
      .master("local[*]")
      .config("spark.default.parallelism", 6)
      .config("spark.sql.shuffle.partitions", 6)
      .getOrCreate()

    rddSortMethod(spark)
    //sparkSQLSortMethod(spark)

    spark.close()
  }

  /**
   * 使用RDD对文件进行排序
   *
   * @param spark
   */
  def rddSortMethod(spark: SparkSession) = {
    val sc = spark.sparkContext
    //读取数据,形成RDD
    val dataRDD: RDD[String] = sc.textFile("data/spark/rdd/order.txt")


    //查看每个分区的数据
    val mapPartitionsWithIndexRDD: RDD[(Int, String)] = dataRDD.mapPartitionsWithIndex {
      case (index, ite) => {
        ite.map(line => {
          (index, line)
        })
      }
    }

    mapPartitionsWithIndexRDD.foreach(println)

    /**
     * 1, 使用sortByKey 完成全局排序
     * sortByKey有2个参数，第一个参数表示是否使用升序排序方式，true表示升序，false表示降序（默认升序）
     * 第二个参数表示分区数，默认不改变分区数，如果传入了分区数，就会按照指定的分区数进行分区
     *
     * 如果直接把时间字段当作key, 得到的结果就是全局有序的的（需要先看0号分区，再看1号分区。。。）
     *
     * 原理：
     * 查看sortByKey源码，可以发现使用的是RangePartitioner
     * 这个分区器，可以保证某一范围内的所有数据都在同一个分区
     * 然后每个分区内，根据key进行排序
     * 这个key的排序规则也是可以指定的
     */

    val sortByKeyRDD = dataRDD.map(line => {
      val jSONObject = JSON.parseObject(line)
      val orderTimeStr = jSONObject.getString("orderTime")

      val orderTime: Date = DateUtils.parseDate(orderTimeStr, "yyyy-MM-dd HH:mm:ss")
      orderTime.getTime

      (orderTime, line)
    }).sortByKey()

    val sortByKeyMapPartitionsWithIndexRDD = sortByKeyRDD.mapPartitionsWithIndex {
      case (index, ite) => {
        ite.map(line => {
          (index, line)
        })
      }
    }

    sortByKeyMapPartitionsWithIndexRDD.foreach(println)

    println("-------下面是使用sortByKey完成分区内排序--------")

    /**
     * 对相同订单号的数据进行排序(分组排序)
     *
     * sortBy 用于指定key
     * SortOrder继承Ordering, 实现排序
     */
    val groupSortRDD: RDD[(SortOrder, String)] = dataRDD.map(line => {
      val jSONObject = JSON.parseObject(line)
      val orderId = jSONObject.getString("orderId")
      val orderTimeStr = jSONObject.getString("orderTime")
      val orderTime = DateUtils.parseDate(orderTimeStr, "yyyy-MM-dd HH:mm:ss").getTime

      val sortOrder = new SortOrder(orderId, orderTime)

      (sortOrder, line)
    }).sortBy(_._1.orderId)

    val groupSortByKeyMapPartitionsWithIndexRDD = groupSortRDD.mapPartitionsWithIndex {
      case (index, ite) => {
        ite.map(t => {
          (index, t._2)
        })
      }
    }

    groupSortByKeyMapPartitionsWithIndexRDD.foreach(println)

    println("-------下面是使用repartitionAndSortWithinPartitions完成分区内排序--------")

    /**
     * repartitionAndSortWithinPartitions
     * 使用指定的分区器进行分区
     * 并在每个分区内根据key进行排序
     */
    dataRDD.map(line => {
      val jSONObject = JSON.parseObject(line)
      val orderId = jSONObject.getString("orderId")
      val orderTimeStr = jSONObject.getString("orderTime")
      val orderTime = DateUtils.parseDate(orderTimeStr, "yyyy-MM-dd HH:mm:ss").getTime
      (orderId+"\001"+orderTime, line)
    }).repartitionAndSortWithinPartitions(new Partitioner {
      override def numPartitions: Int = 2

      override def getPartition(key: Any): Int = {

        val fields = StringUtils.split(key.toString, "\001")

        val orderID = fields(0)

        orderID.hashCode.abs % numPartitions
      }
    }).mapPartitionsWithIndex {
      case (index, ite) => {
        ite.map(t => {
          (index, t._2)
        })
      }
    }.foreach(println)
  }

  class SortOrder(val orderId: String, val orderTime: Long) extends Ordering[SortOrder] with Serializable {
    /**
     * 比较方法，用于排序
     *
     * *  - negative if x < y
     * *  - positive if x > y
     * *  - zero otherwise (if x == y)
     *
     * @param x
     * @param y
     * @return
     */
    override def compare(x: SortOrder, y: SortOrder): Int = {
      var res = 0

      res = if (x.orderTime > y.orderTime) 1 else -1

      res
    }
  }

  /**
   * 使用SparkSQL的方式对文件进行排序
   *
   * @param spark
   */
  def sparkSQLSortMethod(spark: SparkSession) = {

    val dataDF: DataFrame = spark.read.json("data/spark/rdd/order.txt")

    dataDF.createOrReplaceTempView("v_order")

    println("------全局排序--------")
    //全局排序
    spark.sql(
      """
        |SELECT * FROM v_order ORDER BY orderTime
        |""".stripMargin)
      .show()

    println("------分组排序求top10--------")
    //如果不需要求top10，就不需要子查询
    spark.sql(
      """
        | SELECT
        |  orderId,
        |  orderTime
        |FROM (
        |  SELECT
        |    orderId,
        |    orderTime,
        |    row_number() OVER (PARTITION BY orderId ORDER BY orderTime ASC) as rn
        |  FROM v_order) tmp
        |WHERE
        |  rn <= 10
        |""".stripMargin)
      .show()
  }
}
